package org.finesys.chat.core.base.embedding.store.milvus;

import java.lang.reflect.Type;
import java.math.BigDecimal;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.stream.Collectors;

import org.finesys.chat.core.base.document.Metadata;
import org.finesys.chat.core.base.embedding.Embedding;
import org.finesys.chat.core.base.embedding.EmbeddingMatch;
import org.finesys.chat.core.base.embedding.RelevanceScore;
import org.finesys.chat.core.base.embedding.store.util.Generator;
import org.finesys.chat.core.base.segment.TextSegment;
import org.finesys.common.core.util.Utils;

import com.google.common.reflect.TypeToken;
import com.google.gson.Gson;
import com.google.gson.GsonBuilder;
import com.google.gson.JsonObject;
import com.google.gson.ToNumberPolicy;

import io.milvus.client.MilvusServiceClient;
import io.milvus.common.clientenum.ConsistencyLevelEnum;
import io.milvus.exception.IllegalResponseException;
import io.milvus.exception.ParamException;
import io.milvus.response.QueryResultsWrapper;
import io.milvus.response.SearchResultsWrapper;

public class Mapper {

    private static final Gson GSON;
    private static final Type type;

    static {
        GSON = new GsonBuilder().setObjectToNumberStrategy(ToNumberPolicy.LONG_OR_DOUBLE).create();
        type = new TypeToken<Map<String, Object>>() {
        }.getType();
    }


    public static List<List<Double>> toVector(List<Embedding> embeddings) {
        return embeddings.stream().map(Embedding::vectorAsList).collect(Collectors.toList());
    }

    public static List<List<Float>> toVectorFloat(List<Embedding> embeddings) {
        return toVector(embeddings).stream().map((vector) -> vector.stream().map(Double::floatValue).collect(Collectors.toList())).collect(Collectors.toList());
    }

    public static List<String> toText(List<TextSegment> textSegments, int size) {
        return Utils.isNullOrEmpty(textSegments) ? Generator.generateEmptyScalars(size) : toText(textSegments);
    }

    public static List<String> toText(List<TextSegment> textSegments) {
        return textSegments.stream().map(TextSegment::getText).collect(Collectors.toList());
    }

    public static List<JsonObject> toMetadataJsons(List<TextSegment> textSegments, int size) {
        return Utils.isNullOrEmpty(textSegments) ? Generator.generateEmptyJsons(size) : (List<JsonObject>) textSegments.stream().map((segment) -> {
            return GSON.toJsonTree(segment.getMetadata().toMap()).getAsJsonObject();
        }).collect(Collectors.toList());
    }


    public static List<EmbeddingMatch<TextSegment>> toEmbeddingMatches(MilvusServiceClient milvusClient, SearchResultsWrapper searchResultsWrapper, String collectionName, FieldDefinition fieldDefinition, ConsistencyLevelEnum consistencyLevel, boolean queryForVectorOnSearch) {
        List<EmbeddingMatch<TextSegment>> matches = new ArrayList<>();
        Map<String, Embedding> idToEmbedding = new HashMap<>();
        if (queryForVectorOnSearch) {
            try {
                List<String> rowIds = (List<String>) searchResultsWrapper.getFieldWrapper(fieldDefinition.getIdFieldName()).getFieldData();
                idToEmbedding.putAll(queryEmbeddings(milvusClient, collectionName, fieldDefinition, rowIds, consistencyLevel));
            } catch (IllegalResponseException e) {
                throw new RuntimeException(e);
            } catch (ParamException e) {
                throw new RuntimeException(e);
            }
        }

        for (int i = 0; i < searchResultsWrapper.getRowRecords(0).size(); i++) {
            double score = searchResultsWrapper.getIDScore(0).get(i).getScore();
            String rowId = searchResultsWrapper.getIDScore(0).get(i).getStrID();
            Embedding embedding = (Embedding) idToEmbedding.get(rowId);
            TextSegment textSegment = toTextSegment(searchResultsWrapper.getRowRecords(0).get(i), fieldDefinition);
            EmbeddingMatch<TextSegment> embeddingMatch = new EmbeddingMatch<>(RelevanceScore.fromCosineSimilary(score), rowId, embedding, textSegment, i);
            matches.add(embeddingMatch);
        }
        return matches;
    }

    private static TextSegment toTextSegment(QueryResultsWrapper.RowRecord rowRecord, FieldDefinition fieldDefinition) {
        Object textField = rowRecord.get(fieldDefinition.getTextFieldName());
        String text = textField == null ? null : textField.toString();
        if (Utils.isNullOrBlank(text)) {
            return null;
        }
        if (!rowRecord.getFieldValues().containsKey(fieldDefinition.getMetadataFieldName())) {
            return TextSegment.from(text);
        }
        JsonObject jsonObject = (JsonObject) rowRecord.get(fieldDefinition.getMetadataFieldName());
        return TextSegment.from(text, toMetadata(jsonObject));
    }

    private static Metadata toMetadata(JsonObject jsonObject) {
        Map<String, Object> metadataMap = GSON.fromJson(jsonObject, type);
        metadataMap.forEach((key, value) -> {
            if (value instanceof BigDecimal) {
                metadataMap.put(key, ((BigDecimal) value).doubleValue());
            }
        });
        return Metadata.from(metadataMap);
    }

    private static Map<String, Embedding> queryEmbeddings(MilvusServiceClient milvusClient, String collectionName, FieldDefinition fieldDefinition, List<String> rowIds, ConsistencyLevelEnum consistencyLevel) {

        QueryResultsWrapper queryResultsWrapper = CollectionOperationsExecutor.queryForVectors(milvusClient, collectionName, fieldDefinition, rowIds, consistencyLevel);
        Map<String, Embedding> idToEmbedding = new HashMap<>();
        List<QueryResultsWrapper.RowRecord> rowRecords = queryResultsWrapper.getRowRecords();
        rowRecords.stream().forEach((rowRecord) -> {
            String rowId = (String) rowRecord.get(fieldDefinition.getIdFieldName());
            List<Double> vector = (List<Double>) rowRecord.get(fieldDefinition.getVectorFieldName());
            idToEmbedding.put(rowId, Embedding.from(vector));
        });
        return idToEmbedding;
    }
}
